import pickle
from collections import defaultdict
NOTEBOOK_PATH = '/home/fly/notebook/'

def initialize(context):
    run_interval(context, interval_handle, seconds = 3)

def before_trading_start(context, data):
    g.save_flag = False
    g.positions = context.portfolio.positions

    # 设置持久化文件路径
    date_str = context.current_dt.strftime('%Y%m%d')
    file_path = NOTEBOOK_PATH + f'buy_order_record{date_str}.pkl'

    # 现价偏离下单价的幅度
    g.order_adjust_percentile = 1.005
    # 触发上涨直接下单的涨幅
    g.up_thredhold = 5.5
    # 触发上涨策略下单的涨幅
    g.up_observe_thredhold = 2.5
    
    # 获取资产价值和可用现金
    g.portfolio_value = context.portfolio.portfolio_value
    g.cash = context.portfolio.portfolio_value
    # 个股的持仓比例
    g.position_ratio = 0.1
    g.position_unit_value = g.portfolio_value * g.position_ratio
    g.stock_buy_list = {"600570.SS": {"buy_thredhold": 3}}

    try:
        with open(file_path, 'rb') as f:
            g.order_dict = pickle.load(f)
    except:
        g.order_dict = defaultdict(list)

def place_buy_order(sid, sid_snapshot, positions):

    # 获取买一价格和当前价格
    buy1_price = sid_snapshot[sid]['bid_grp'][1][0]
    close_price = sid_snapshot[sid]['close_price']

    sid_target_buy_value = g.position_unit_value
    # 如果当前已经有该股持仓，且持仓市值超过设定的单位阈值，则不进行交易
    if sid in positions and positions.amount*close_price >= g.position_unit_value:
        return
    # 如果当前已经有该股持仓，且持仓市值超过设定的单位阈值，则计算需要补仓的金额
    elif sid in positions and positions.amount*close_price < g.position_unit_value:
        sid_target_buy_value = g.position_unit_value - positions.amount*close_price

    # 如果相应的股票ID已经下过单
    if sid in g.order_dict:
        # 如果股票ID的订单状态为8（已成）或者10（不需要重复下单）则直接返回
        if g.order_dict[sid]["status"] in ["8", "10"]:
            return
        # 如果股票ID的订单状态不等于8（已成）或者10（不需要重复下单）
        elif g.order_dict[sid]["status"] not in ["8", "10"]:
            # 获取最新的订单信息
            order_id = g.order_dict[sid]["order_id"]
            order_info = get_order(order_id)
            order_status = order_info['status']
            # 如果股票ID的订单状态为8（已成），更新订单状态和持仓信息
            if order_status == "8":
                g.order_dict[sid] = order_info
                g.positions = context.portfolio.positions
                g.save_flag = True
            # 如果股票ID的订单状态不是8（已成），结合交易信息和订单信息判断是否重新下单
            else:
                order_amount = order_info['amount']
                order_filled = order_info['filled']
                limit_price = order_info['limit']
                
                # 当下单数量不等于已成交数量时
                if order_amount != order_filled:
                    # 计算未成交金额是否大于1手的股票价格
                    remain_target_value = abs(order_amount - order_filled) * limit_price
                    remain_amount = remain_target_value // close_price * 100
                    # 未成交金额小于1手的股票价格，取消订单，更新股票下单信息为10（不需要重复下单）
                    if remain_amount == 0:
                        cancle_order(order_id)
                        g.order_dict[sid]["status"] = "10"
                        g.save_flag = True
                    
                    # 未成交金额大于1手的股票价格，且当前价格变化幅度超过之前的下单阈值，则取消订单，并以最新价格下单
                    elif remain_amount != 0 and close_price >= g.order_adjust_percentile*limit_price:
                        cancle_order(order_id)
                        place_order_action(sid, remain_amount, buy1_price)
    # 如果相应的股票ID没有下过单
    elif sid not in g.order_dict:
        # 计算下单数量
        g.cash = context.portfolio.portfolio_value
        target_amount = min(g.cash, sid_target_buy_value) // close_price * 100

        # 当下单数量大于等1手的股票价格，则下单买入，否则不错操作
        if target_amount != 0:
            place_order_action(sid, target_amount, limit_price)

    # 如果有订单状态发生变更，则保存订单状态到本地文件，避免策略宕机
    if g.save_flag:
        with open(file_path, 'wb') as f:
            pickle.dump(g.order_dict, f, -1)
            g.save_flag = False

def place_order_action(sid, place_amount, limit_price):
    '''
        执行下单操作
    '''
    order_id = order(sid, place_amount, limit_price)
    order_info = get_order(order_id)
    g.order_dict[sid] = order_info
    g.save_flag = True

def signal_moving_average(sid_price_df, short_period, long_period):
    
    # 计算5分钟和21分钟的均线（确保从有足够数据点的索引开始）
    sid_price_df['5min_avg'] = sid_price_df['close'].rolling(window=short_period, min_periods=short_period).mean()
    sid_price_df['21min_avg'] = sid_price_df['close'].rolling(window=long_period, min_periods=long_period).mean()

    # 初始化买卖信号列并填充为0（从有足够数据点的索引开始）
    sid_price_df['buy_signal'] = 0
    sid_price_df['sell_signal'] = 0


    # 找出买卖信号（从有足够数据点的索引开始）
    # 实际上这里取5，因为需要5个数据点来计算5分钟均线
    start_idx = max(short_period - 1, long_period - 1)  

    for i in range(start_idx, len(sid_price_df)):
        # 检查前一个值是否为NaN（即确保有足够的数据点来计算均线）
        if not (pd.isnull(sid_price_df.at[i - 1, '5min_avg']) or pd.isnull(sid_price_df.at[i - 1, '21min_avg'])):
            # 卖出信号
            if sid_price_df.iloc[i - 1]['5min_avg'] > sid_price_df.iloc[i - 1]['21min_avg'] and sid_price_df.iloc[i]['5min_avg'] < sid_price_df.iloc[i]['21min_avg']:
                sid_price_df.loc[sid_price_df.index[i], 'sell_signal'] = 1
            # 买入信号
            if sid_price_df.iloc[i - 1]['5min_avg'] < sid_price_df.iloc[i - 1]['21min_avg'] and sid_price_df.iloc[i]['5min_avg'] > sid_price_df.iloc[i]['21min_avg']:
                sid_price_df.loc[sid_price_df.index[i], 'buy_signal'] = 1
        
    return sid_price_df.iloc[-1]['buy_signal'], sid_price_df.iloc[-1]['sell_signal']



def trading_signal(sid, start_time, end_time):




def interval_handle(context):
    sid = g.security
    start_time = context.current_dt.strftime('%Y%m%d') + "0930"
    
    # 遍历待买入列表的股票
    for sid in g.stock_buy_list:
        sid_snapshot = get_snapshot(sid)
        px_change_rate = sid_snapshot['px_change_rate']

        # 当涨幅超过预设阈值，则直接买入
        if px_change_rate >= g.up_thredhold:
            place_buy_order(sid, sid_snapshot, g.positions)
        # 当涨跌幅没有超过预设阈值，则进行策略下单
        else:
            # 获取当前时间
            end_time = context.current_dt.strftime('%Y%m%d%H%M')        
            # 获取行情价格
            sid_price_df = get_price(sid, start_time, end_time, frequency = '1m', fields = 'close')
    
            # 设置长短周期
            short_period = 5
            long_period = 21
            
            # 获取交易信号
            buy_signal, sell_signal = signal_moving_average(sid_price_df, short_period, long_period)

            if buy_signal == 1 and px_change_rate >= g.up_observe_thredhold:
                place_buy_order(sid, sid_snapshot, g.positions)

def handle_data(context, data):
    pass